package com.renne.ai.dbcmcp.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.renne.ai.dbcmcp.service.ChatService;
import com.renne.ai.dbcmcp.service.DBTableStructureService;
import com.renne.ai.dbcmcp.service.RolePromptService;
import io.modelcontextprotocol.server.McpServer;
import io.modelcontextprotocol.server.McpServerFeatures;
import io.modelcontextprotocol.server.McpSyncServer;
import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.server.transport.StdioServerTransportProvider;
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
import io.modelcontextprotocol.spec.McpSchema;
import io.modelcontextprotocol.spec.McpServerTransportProvider;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.mcp.McpToolUtils;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.ToolCallbacks;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.reactive.function.server.RouterFunction;
import org.springframework.web.server.WebExceptionHandler;
import reactor.core.publisher.Mono;

import javax.annotation.Resource;
import java.util.List;
import java.util.function.BiConsumer;
@Slf4j
@Configuration
public class McpServerConfig {

    @Resource
    private DBTableStructureService dbTableStructureService;
    @Resource
    private RolePromptService rolePromptService;
    @Resource
    private ChatService chatService;


    // STDIO transport
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "stdio")
    public StdioServerTransportProvider stdioServerTransportProvider() {
        return new StdioServerTransportProvider();
    }

    // SSE transport
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
    public WebFluxSseServerTransportProvider sseServerTransportProvider() {
        return new WebFluxSseServerTransportProvider(new ObjectMapper(), "/mcp/message");
    }

    // Router function for SSE transport used by Spring WebFlux to start an HTTP
    // server.
    // 用于SSE传输的路由函数，由Spring WebFlux使用来启动HTTP服务器
    @Bean
    @ConditionalOnProperty(prefix = "transport", name = "mode", havingValue = "sse")
    public RouterFunction<?> mcpRouterFunction(WebFluxSseServerTransportProvider transportProvider) {
        return transportProvider.getRouterFunction();
    }


    @Bean
    public McpSyncServer mcpServer(McpServerTransportProvider transportProvider) {

        // Configure server capabilities with resource support
        var capabilities = McpSchema.ServerCapabilities.builder()
                .tools(true) // Tool support with list changes notifications
                .logging() // Logging support
                .build();

        // 从 dbTableStructureService 中提取 ToolCallback
        ToolCallback[] toolCallbacks = ToolCallbacks.from(dbTableStructureService, rolePromptService,chatService);

        // 将 ToolCallback 转换为 SyncToolSpecification
        List<McpServerFeatures.SyncToolSpecification> syncToolSpecifications = McpToolUtils.toSyncToolSpecifications(toolCallbacks);

        // Create the server with both tool and resource capabilities
        McpSyncServer server = McpServer.sync(transportProvider)
                .serverInfo("DBC MCP Server", "0.0.1")
                .capabilities(capabilities)
                .tools(syncToolSpecifications) // 使用转换后的 toolSpecifications
                .build();

        return server;
    }

    // 日志处理器，捕获异常并记录堆栈
    @Bean
    public BiConsumer<McpSyncServerExchange, String> logHandler() {

        return (exchange, logMessage) -> {
            try {
                log.info("MCP event: {}", logMessage);
                // 模拟可能抛出异常的逻辑
                if (logMessage.contains("error")) {
                    throw new RuntimeException("Simulated MCP error");
                }
            } catch (Throwable t) {
                log.error("MCP event processing failed: {}", logMessage, t); // 记录堆栈
            }
        };
    }
    @Bean
    public WebExceptionHandler customExceptionHandler() {
        return (exchange, ex) -> {
            ex.printStackTrace(); // 打印堆栈信息
            return Mono.error(ex);
        };
    }

}
